# -*- coding: utf-8 -*-
"""
Calibration for zurich instruments
Copyright (C) 2020 Ziyu Tao
"""

import numpy as np
from scipy import optimize
import textwrap
import time


_zurich_delay = 0.06


def spectDeInit(spec):
    spec.gpib_write(':INIT:CONT ON')


def spectFreq(spec, freq):
    spec.gpib_write(':FREQ:CENT %gGHz' % freq)


def spec_init(spec, reg):
    spec.select_device(reg['spec_id'])
    spec.gpib_write(":FORM INT,32:FORM:BORD NORM")
    spec.gpib_write(textwrap.dedent("""\
*RST;
:FREQ:SPAN 100Hz;\
"""))
#     spec.gpib_write(textwrap.dedent("""\
# *RST;:POW:RF:ATT 0dB;\
# :AVER:STAT OFF;:FREQ:SPAN 100Hz;\
# :BAND 300Hz;:INIT:CONT OFF;\
# :SYST:PORT:IFVS:ENAB OFF;\
# :SENS:SWE:POIN 101
# """))


def signalPower(spec):
    """returns the mean power in mW read by the
    spectrum analyzer"""
    # spec.gpib_write('*trg;*opc?;')
    dBsRaw = spec.gpib_query(":TRAC? TRACE1")
    dBs = np.asarray(eval(dBsRaw))
    res = 10.0**(0.1*np.mean(dBs))
    return res


def measurePower(spec, fpga, a, b):
    """returns signal power from the spectrum analyzer
    Args:
        spec: server of spectrum analyzer
        fpga (object): server of zurich device
    """
    fpga.dc_offset(dc_I=a, dc_Q=b)
    time.sleep(_zurich_delay)
    return signalPower(spec)


def minPos(yl, y0, yr):
    """Calculates minimum of a parabola to
    three equally spaced points.
    The return value is in units of the spacing
    relative to the center point.
    It is bounded by -1 and 1.
    Example:
    f(x) =ax^2+bx+c, x_center = -b/(2a)
    yl,y0,yr = f(-1),f(0),f(1)
    0.5*(yl-yr)/(yl+yr-2.0*ym) = -b/(2a)
    """
    x_center = yl+yr-2.0*y0
    if x_center <= 0:
        return 0
    x_center = 0.5*(yl-yr)/x_center
    if x_center > 1:
        x_center = 1
    elif x_center < -1:
        x_center = -1
    return x_center


def scan_spectrum(cxn, reg, f_center=5.0, f_span=0.4):
    """get the traces data from spectrum_analyzer
    """
    spec = cxn.spectrum_analyzer_server
    spec.select_device(reg['spec_id'])
    spec.gpib_write(":FORM INT,32:FORM:BORD NORM")

    spec.gpib_write(':FREQ:CENT %gGHz' % f_center)
    spec.gpib_write(':FREQ:SPAN %gGHz' % f_span)

    dBsRaw = spec.gpib_query(":TRAC? TRACE1")
    dBs = np.asarray(eval(dBsRaw))

    freqs = np.linspace(f_center-f_span/2., f_center+f_span/2., len(dBs))
    return (freqs, dBs)


# -----zero calibration-----

def zero(LO_Server, spec, fpga, freq, noisy=False):
    """Calibrates the zeros for the two ports A abd B of AWG
    using the spectrum analyzer.
    Args:
        freq (float): local frequency of a microwave source
        in GHz.
    """
    LO_Server.frequency(freq*1e3)
    spectFreq(spec, freq)

    a = 0
    b = 0
    max_offset = 1
    precision = 1 << 15
    precision0 = 1 << 15

    while precision > 0:
        step = max_offset*precision/precision0
        al = measurePower(spec, fpga, a-step, b)
        ar = measurePower(spec, fpga, a+step, b)
        ac = measurePower(spec, fpga, a, b)
        corra = int(round(precision*minPos(al, ac, ar)))
        a += corra*max_offset/precision0

        bl = measurePower(spec, fpga, a, b-step)
        br = measurePower(spec, fpga, a, b+step)
        bc = measurePower(spec, fpga, a, b)
        corrb = int(round(precision*minPos(bl, bc, br)))
        b += corrb*max_offset/precision0

        optprec = 2*np.max([np.abs(corra), np.abs(corrb)])
        precision = np.min([optprec, precision // 2.0])
        minpower = 10 * np.log10(bc)
        if noisy:
            print('a = %.5f  b = %.5f uncertainty : %.3f, power %6.1f dBm' %
                  (a, b, precision, minpower))
    return [a, b, minpower]


def zeroFixedCarrier(cxn, LO_Server, reg, fpga, freq, noisy=False):
    # import labrad; cxn = labrad.connect()
    # reg = keys.new_registry(cxn, board_name='dev8334', use_default=True)

    spec = cxn.spectrum_analyzer_server
    spec_init(spec, reg)

    LO_ID = reg['LO_id']
    LO_power = reg['LO_power']

    LO_Server.select_device(LO_ID)
    LO_Server.amplitude(LO_power)
    LO_Server.output(True)

    daczeros = zero(LO_Server, spec, fpga, freq, noisy)
    # LO_Server.output(False)
    spectDeInit(spec)
    return daczeros


# -----sideband calibration-----

def measureOppositeSideband(
        spec, fpga,
        carrierfreq, sidebandfreq, amp_sb, compensation):
    """
    Args:
        compensation (list or tuple): [delta_amp, delta_phi]
    """

    amp0 = amp_sb
    delta_amp, delta_phi = compensation

    fpga.sine_amp(amp0 + delta_amp, amp0 - delta_amp)
    fpga.shift_phase(0. + delta_phi, -90. - delta_phi)
    fpga.sine_freq(sidebandfreq*1e9)
    fpga.open_sine_pulse()
    time.sleep(_zurich_delay)

    dBsRaw = spec.gpib_query(":TRAC? TRACE1")
    dBs = np.asarray(eval(dBsRaw))
    return np.mean(dBs)


def _sideband(LO_Server, spec, fpga, carrierfreq, sidebandfreq,
              amp_sb=0.2, noisy=False, close_fpga=False, method='Powell',
              x0=None, ratio_amp=57.3, max_step=50):
    """When the IQ mixer is used for arbitrary wave generator,
    imperfections in the IQ mixer and the DACs give rise to
    a signal at both
    carrierfreq + sidebandfreq (target part)
    carrierfreq - sidebandfreq (mirror part).

    This function determines amplitude and phase compensation of
    the sideband signal for zurich instrument, to inhibit the
    mirror part.
    Example:
    fpga.sine_amp(amp0 + x[0], amp0 - x[0])
    fpga.shift_phase(0. + x[1], -90. - x[1])

    Args:
        method (str): if you do not have better method,
        just use method = 'Powell' !!
        ratio_amp (float): determined by x[1]/x[0], with
        the value around 180/pi ~ 57.3, since zurich use angle unit
        but not rad unit.
    """
    LO_Server.frequency(carrierfreq*1e3)
    spectFreq(spec, carrierfreq-sidebandfreq)

    def wrapped_minimize(x):
        delta_amp = x[0]/ratio_amp
        delta_phi = x[1]
        power = measureOppositeSideband(
            spec, fpga, carrierfreq, sidebandfreq, amp_sb,
            compensation=[delta_amp, delta_phi])
        if noisy:
            print(x, power)
        return power

    if x0 is None:
        x0 = [0., 0.]

    options = {'maxfev': max_step, 'xtol': 10.}
    res = optimize.minimize(
        wrapped_minimize, x0=x0, method=method,
        options=options)
    delta_amp = res.x[0]/ratio_amp
    delta_phi = res.x[1]
    minpower_dB = res.fun

    return np.hstack([delta_amp, delta_phi, minpower_dB])


def sideband(
        cxn, LO_Server, reg, fpga, carrierfreq, sidebandfreq,
        x0=None, amp_sb=0.2, ratio_amp=None, max_step=30,
        accept_power=-70, recursion=3):
    """initiate devices and run _sideband
    """
    if ratio_amp is None:
        ratio_amp = 180./np.pi
    spec = cxn.spectrum_analyzer_server
    spec_init(spec, reg)

    LO_ID = reg['LO_id']
    LO_power = reg['LO_power']
    LO_Server.select_device(LO_ID)
    LO_Server.amplitude(LO_power)
    LO_Server.output(True)

    data = _sideband(LO_Server, spec, fpga, carrierfreq, sidebandfreq,
                     amp_sb=amp_sb, noisy=False, close_fpga=False,
                     method='Powell', x0=x0, ratio_amp=ratio_amp,
                     max_step=max_step)

    minpower_dB = data[-1]

    do_return = (minpower_dB < accept_power) or (recursion - 1 <= 0)
    if do_return:
        return data
    else:
        x0_update = [data[0]*ratio_amp, data[1]]
        data = sideband(
            cxn, LO_Server, reg, fpga, carrierfreq, sidebandfreq,
            x0=x0_update, amp_sb=amp_sb, ratio_amp=ratio_amp,
            max_step=max_step, recursion=recursion-1)
        return data
